package socket

import (
	"bufio"
	"errors"
	"log/slog"
	"net"
	"reflect"
)

const (
	UDP byte = 0
	TCP byte = 1
)

type Socket struct {
	socketType byte
	Reader     *bufio.Reader
	Writer     *bufio.Writer
	conn       net.Conn
	udpWrap    *UdpWrap
}

func NewSocket(conn net.Conn) *Socket {
	return &Socket{socketType: TCP, Reader: bufio.NewReader(conn), Writer: bufio.NewWriter(conn), conn: conn}
}

func NewSocketUDP(conn net.Conn, udpWrap *UdpWrap) *Socket {
	return &Socket{
		socketType: UDP,
		udpWrap:    udpWrap,
		Reader:     bufio.NewReader(udpWrap),
		Writer:     bufio.NewWriter(udpWrap),
		conn:       conn,
	}
}

func (p *Socket) LocalAddr() net.Addr {
	if p.conn != nil {
		return p.conn.LocalAddr()
	}
	return nil
}

func (p *Socket) RemoteAddr() net.Addr {
	if p.socketType == UDP && p.udpWrap != nil {
		return p.udpWrap.RemoteAddr()
	} else if p.conn != nil {
		return p.conn.RemoteAddr()
	}
	return nil
}

func (p *Socket) Wait() error {
	if p.conn == nil || p.Reader == nil {
		return errors.New("socket connection not found")
	}

	if p.Reader.Buffered() > 0 {
		return nil
	}
	if _, err := p.Reader.Peek(1); err != nil {
		return err
	}
	return nil
}

func (p *Socket) Read(buf []byte) (n int, err error) {
	return p.Reader.Read(buf)
}

func (p *Socket) ReadByte() (byte, error) {
	return p.Reader.ReadByte()
}

func (p *Socket) ReaderBytes(bytes []byte) (int, error) {
	return p.Reader.Read(bytes)
}

func (p *Socket) Write(buf []byte) (n int, err error) {
	return p.Writer.Write(buf)
}

func (p *Socket) WriteByte(b byte) error {
	return p.Writer.WriteByte(b)
}

func (p *Socket) WriterBytes(bytes []byte) error {
	offset := 0
	for {
		if nn, err := p.Writer.Write(bytes[offset:]); err != nil {
			_ = p.Writer.Flush()
			return err
		} else {
			offset += nn
		}
		if offset == len(bytes) {
			if err := p.Writer.Flush(); err != nil {
				return err
			}
			return nil
		}
	}
}

func (p *Socket) Flush() error {
	return p.Writer.Flush()
}

func (p *Socket) Close() error {
	if p.conn != nil {
		slog.Debug(reflect.ValueOf(p.conn).String()+" Close connection", "local", p.LocalAddr(), "remote", p.RemoteAddr())
		return p.conn.Close()
	}
	return nil
}

func (p *Socket) Conn() net.Conn {
	return p.conn
}

func (p *Socket) IsUdp() bool {
	return p.socketType == UDP
}

func (p *Socket) IsTcp() bool {
	return p.socketType == TCP
}

func (p *Socket) UdpWrap() *UdpWrap {
	return p.udpWrap
}
